{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tFhqrPopctwa"
      },
      "source": [
        "Copyright 2020 DeepMind Technologies Limited.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zuYtT7GpuVgQ"
      },
      "source": [
        "\u003cp align=\"center\"\u003e\n",
        "  \u003ch1 align=\"center\"\u003eTAPIR: Tracking Any Point with per-frame Initialization and temporal Refinement\u003c/h1\u003e\n",
        "  \u003cp align=\"center\"\u003e\n",
        "    \u003ca href=\"http://www.carldoersch.com/\"\u003eCarl Doersch\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://yangyi02.github.io/\"\u003eYi Yang\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://scholar.google.com/citations?user=Jvi_XPAAAAAJ\"\u003eMel Vecerik\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://scholar.google.com/citations?user=cnbENAEAAAAJ\"\u003eDilara Gokay\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://www.robots.ox.ac.uk/~ankush/\"\u003eAnkush Gupta\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"http://people.csail.mit.edu/yusuf/\"\u003eYusuf Aytar\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://scholar.google.co.uk/citations?user=IUZ-7_cAAAAJ\"\u003eJoao Carreira\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://www.robots.ox.ac.uk/~az/\"\u003eAndrew Zisserman\u003c/a\u003e\n",
        "  \u003c/p\u003e\n",
        "  \u003ch3 align=\"center\"\u003e\u003ca href=\"https://arxiv.org/abs/2306.08637\"\u003ePaper\u003c/a\u003e | \u003ca href=\"https://deepmind-tapir.github.io\"\u003eProject Page\u003c/a\u003e | \u003ca href=\"https://github.com/deepmind/tapnet\"\u003eGitHub\u003c/a\u003e | \u003ca href=\"https://github.com/deepmind/tapnet/tree/main#running-tapir-locally\"\u003eLive Demo\u003c/a\u003e \u003c/h3\u003e\n",
        "  \u003cdiv align=\"center\"\u003e\u003c/div\u003e\n",
        "\u003c/p\u003e\n",
        "\n",
        "\u003cp align=\"center\"\u003e\n",
        "  \u003ca href=\"\"\u003e\n",
        "    \u003cimg src=\"https://storage.googleapis.com/dm-tapnet/swaying_gif.gif\" alt=\"Logo\" width=\"50%\"\u003e\n",
        "  \u003c/a\u003e\n",
        "\u003c/p\u003e\n",
        "\n",
        ""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mCmDvfFvxnGB"
      },
      "outputs": [],
      "source": [
        "# @title Install code and dependencies {form-width: \"25%\"}\n",
        "!pip install git+https://github.com/google-deepmind/tapnet.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zyEo9-Kv78S7"
      },
      "outputs": [],
      "source": [
        "MODEL_TYPE = 'bootstapir' # 'tapir' or 'bootstapir'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HaswJZMq9B3c"
      },
      "outputs": [],
      "source": [
        "# @title Download Model {form-width: \"25%\"}\n",
        "\n",
        "%mkdir tapnet/checkpoints\n",
        "\n",
        "if MODEL_TYPE == 'tapir':\n",
        "  !wget -P tapnet/checkpoints https://storage.googleapis.com/dm-tapnet/causal_tapir_checkpoint.npy\n",
        "else:\n",
        "  !wget -P tapnet/checkpoints https://storage.googleapis.com/dm-tapnet/bootstap/causal_bootstapir_checkpoint.npy\n",
        "\n",
        "%ls tapnet/checkpoints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FxlHY242m-6Q"
      },
      "outputs": [],
      "source": [
        "# @title Imports {form-width: \"25%\"}\n",
        "%matplotlib widget\n",
        "import functools\n",
        "\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "import mediapy as media\n",
        "import numpy as np\n",
        "from tqdm import tqdm\n",
        "import tree\n",
        "\n",
        "from tapnet.models import tapir_model\n",
        "from tapnet.utils import transforms\n",
        "from tapnet.utils import viz_utils\n",
        "from tapnet.utils import model_utils\n",
        "\n",
        "from google.colab import output\n",
        "output.enable_custom_widget_manager()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7rfy2yobnHqw"
      },
      "outputs": [],
      "source": [
        "# @title Load Checkpoint {form-width: \"25%\"}\n",
        "\n",
        "if MODEL_TYPE == 'tapir':\n",
        "  checkpoint_path = 'tapnet/checkpoints/causal_tapir_checkpoint_panning.npy'\n",
        "else:\n",
        "  checkpoint_path = 'tapnet/checkpoints/causal_bootstapir_checkpoint.npy'\n",
        "ckpt_state = np.load(checkpoint_path, allow_pickle=True).item()\n",
        "params, state = ckpt_state['params'], ckpt_state['state']\n",
        "\n",
        "kwargs = dict(use_causal_conv=True, bilinear_interp_with_depthwise_conv=False, pyramid_level=0)\n",
        "if MODEL_TYPE == 'bootstapir':\n",
        "  kwargs.update(dict(\n",
        "    pyramid_level=1,\n",
        "    extra_convs=True,\n",
        "    softmax_temperature=10.0\n",
        "  ))\n",
        "tapir = tapir_model.ParameterizedTAPIR(params, state, tapir_kwargs=kwargs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I7wOMJoSQzq1"
      },
      "outputs": [],
      "source": [
        "# @title Build Model {form-width: \"25%\"}\n",
        "\n",
        "# Internally, the tapir model has three stages of processing: computing\n",
        "# image features (get_feature_grids), extracting features for each query point\n",
        "# (get_query_features), and estimating trajectories given query features and\n",
        "# the feature grids where we want to track (estimate_trajectories).  For\n",
        "# tracking online, we need extract query features on the first frame only, and\n",
        "# then call estimate_trajectories on one frame at a time.\n",
        "\n",
        "def online_model_init(frames, query_points):\n",
        "  \"\"\"Initialize query features for the query points.\"\"\"\n",
        "  frames = model_utils.preprocess_frames(frames)\n",
        "  feature_grids = tapir.get_feature_grids(frames, is_training=False)\n",
        "  query_features = tapir.get_query_features(\n",
        "      frames,\n",
        "      is_training=False,\n",
        "      query_points=query_points,\n",
        "      feature_grids=feature_grids,\n",
        "  )\n",
        "  return query_features\n",
        "\n",
        "online_model_init=jax.jit(online_model_init)\n",
        "\n",
        "def online_model_predict(frames, query_features, causal_context):\n",
        "  \"\"\"Compute point tracks and occlusions given frames and query points.\"\"\"\n",
        "  frames = model_utils.preprocess_frames(frames)\n",
        "  feature_grids = tapir.get_feature_grids(frames, is_training=False)\n",
        "  trajectories = tapir.estimate_trajectories(\n",
        "      frames.shape[-3:-1],\n",
        "      is_training=False,\n",
        "      feature_grids=feature_grids,\n",
        "      query_features=query_features,\n",
        "      query_points_in_video=None,\n",
        "      query_chunk_size=64,\n",
        "      causal_context=causal_context,\n",
        "      get_causal_context=True,\n",
        "  )\n",
        "  causal_context = trajectories['causal_context']\n",
        "  del trajectories['causal_context']\n",
        "  # Take only the predictions for the final resolution.\n",
        "  # For running on higher resolution, it's typically better to average across\n",
        "  # resolutions.\n",
        "  tracks = trajectories['tracks'][-1]\n",
        "  occlusions = trajectories['occlusion'][-1]\n",
        "  uncertainty = trajectories['expected_dist'][-1]\n",
        "  visibles = model_utils.postprocess_occlusions(occlusions, uncertainty)\n",
        "  return tracks, visibles, causal_context\n",
        "\n",
        "online_model_predict=jax.jit(online_model_predict)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ogRTRVgfSq0W"
      },
      "outputs": [],
      "source": [
        "# @title Utility Functions {form-width: \"25%\"}\n",
        "\n",
        "def sample_random_points(frame_max_idx, height, width, num_points):\n",
        "  \"\"\"Sample random points with (time, height, width) order.\"\"\"\n",
        "  y = np.random.randint(0, height, (num_points, 1))\n",
        "  x = np.random.randint(0, width, (num_points, 1))\n",
        "  t = np.random.randint(0, frame_max_idx + 1, (num_points, 1))\n",
        "  points = np.concatenate((t, y, x), axis=-1).astype(np.int32)  # [num_points, 3]\n",
        "  return points"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Izp33JBg6eij"
      },
      "outputs": [],
      "source": [
        "# @title Load an Exemplar Video {form-width: \"25%\"}\n",
        "\n",
        "%mkdir tapnet/examplar_videos\n",
        "\n",
        "!wget -P tapnet/examplar_videos https://storage.googleapis.com/dm-tapnet/horsejump-high.mp4\n",
        "\n",
        "video = media.read_video('tapnet/examplar_videos/horsejump-high.mp4')\n",
        "height, width = video.shape[1:3]\n",
        "media.show_video(video, fps=10)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_LLK7myqp3Px"
      },
      "outputs": [],
      "source": [
        "# @title Progressively Predict Sparse Point Tracks {form-width: \"25%\"}\n",
        "\n",
        "resize_height = 256  # @param {type: \"integer\"}\n",
        "resize_width = 256  # @param {type: \"integer\"}\n",
        "num_points = 20  # @param {type: \"integer\"}\n",
        "\n",
        "frames = media.resize_video(video, (resize_height, resize_width))\n",
        "query_points = sample_random_points(0, frames.shape[1], frames.shape[2], num_points)\n",
        "\n",
        "# Extract features for the query point.\n",
        "query_features = online_model_init(frames[None, 0:1], query_points[None])\n",
        "\n",
        "# We need a causal state for every PIPs resolution; therefore,\n",
        "# construct_initial_causal_state's second param is the number of resolutions.\n",
        "# query_features.resolutions is one more than the number of PIPs resolutions\n",
        "# (since there's an 'initialization' resolution as well).\n",
        "# Of course, since we're running at 256x256 in this demo, there will only be\n",
        "# one PIPs resolution.\n",
        "#\n",
        "# Note that causal_state can also be initialized to None, but this will mean\n",
        "# online_model_predict gets compiled twice.\n",
        "causal_state = tapir.construct_initial_causal_state(query_points.shape[0], len(query_features.resolutions) - 1)\n",
        "\n",
        "# Predict point tracks frame by frame\n",
        "predictions = []\n",
        "for i in range(frames.shape[0]):\n",
        "  # Note: we add a batch dimension.\n",
        "  tracks, visibles, causal_state = online_model_predict(\n",
        "      frames=frames[None, i:i+1],\n",
        "      query_features=query_features,\n",
        "      causal_context=causal_state,\n",
        "  )\n",
        "  predictions.append({'tracks':tracks, 'visibles':visibles})\n",
        "\n",
        "tracks = np.concatenate([x['tracks'][0] for x in predictions], axis=1)\n",
        "visibles = np.concatenate([x['visibles'][0] for x in predictions], axis=1)\n",
        "\n",
        "# Visualize sparse point tracks\n",
        "tracks = transforms.convert_grid_coordinates(tracks, (resize_width, resize_height), (width, height))\n",
        "video_viz = viz_utils.plot_tracks_v2(video, tracks, np.logical_not(visibles))\n",
        "media.show_video(video_viz, fps=10)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9RyQhPLrwpjF"
      },
      "outputs": [],
      "source": [
        "# @title Select Any Points at First Frame {form-width: \"25%\"}\n",
        "\n",
        "# Generate a colormap with 20 points, no need to change unless select more than 20 points\n",
        "colormap = viz_utils.get_colors(20)\n",
        "\n",
        "fig, ax = plt.subplots(figsize=(10, 5))\n",
        "ax.imshow(video[0])\n",
        "ax.axis('off')\n",
        "ax.set_title('You can select more than 1 points. After select enough points, run the next cell.')\n",
        "\n",
        "select_points = []\n",
        "\n",
        " # Event handler for mouse clicks\n",
        "def on_click(event):\n",
        "  if event.button == 1 and event.inaxes == ax:  # Left mouse button clicked\n",
        "    x, y = int(np.round(event.xdata)), int(np.round(event.ydata))\n",
        "\n",
        "    select_points.append(np.array([x, y]))\n",
        "\n",
        "    color = colormap[len(select_points) - 1]\n",
        "    color = tuple(np.array(color) / 255.0)\n",
        "    ax.plot(x, y, 'o', color=color, markersize=5)\n",
        "    plt.draw()\n",
        "\n",
        "fig.canvas.mpl_connect('button_press_event', on_click)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J5GVZDywwsO-"
      },
      "outputs": [],
      "source": [
        "# @title Predict Point Tracks for the Selected Points {form-width: \"25%\"}\n",
        "\n",
        "resize_height = 256  # @param {type: \"integer\"}\n",
        "resize_width = 256  # @param {type: \"integer\"}\n",
        "\n",
        "def convert_select_points_to_query_points(frame, points):\n",
        "  \"\"\"Convert select points to query points.\n",
        "\n",
        "  Args:\n",
        "    points: [num_points, 2], [t, y, x]\n",
        "  Returns:\n",
        "    query_points: [num_points, 3], [t, y, x]\n",
        "  \"\"\"\n",
        "  points = np.stack(points)\n",
        "  query_points = np.zeros(shape=(points.shape[0], 3), dtype=np.float32)\n",
        "  query_points[:, 0] = frame\n",
        "  query_points[:, 1] = points[:, 1]\n",
        "  query_points[:, 2] = points[:, 0]\n",
        "  return query_points\n",
        "\n",
        "frames = media.resize_video(video, (resize_height, resize_width))\n",
        "query_points = convert_select_points_to_query_points(0, select_points)\n",
        "query_points = transforms.convert_grid_coordinates(\n",
        "    query_points, (1, height, width), (1, resize_height, resize_width), coordinate_format='tyx')\n",
        "\n",
        "query_features = online_model_init(frames[None, 0:1], query_points[None])\n",
        "causal_state = tapir.construct_initial_causal_state(query_points.shape[0], len(query_features.resolutions) - 1)\n",
        "\n",
        "# Predict point tracks frame by frame\n",
        "predictions = []\n",
        "for i in range(frames.shape[0]):\n",
        "  # Note: we add a batch dimension.\n",
        "  tracks, visibles, causal_state = online_model_predict(\n",
        "      frames=frames[None, i:i+1],\n",
        "      query_features=query_features,\n",
        "      causal_context=causal_state,\n",
        "  )\n",
        "  predictions.append({'tracks':tracks, 'visibles':visibles})\n",
        "\n",
        "tracks = np.concatenate([x['tracks'][0] for x in predictions], axis=1)\n",
        "visibles = np.concatenate([x['visibles'][0] for x in predictions], axis=1)\n",
        "\n",
        "# Visualize sparse point tracks\n",
        "tracks = transforms.convert_grid_coordinates(tracks, (resize_width, resize_height), (width, height))\n",
        "video_viz = viz_utils.paint_point_track(video, tracks, visibles, colormap)\n",
        "media.show_video(video_viz, fps=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8orfQRoaRJit"
      },
      "source": [
        "That's it!"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
